package tester

import (
	"encoding/hex"
	"fmt"
	"log"
	"math/big"
	"math/rand"
	"net"
	"strconv"
	"time"
)

func genTranId() string {
	hexString := "0123456789abcdef"
	s := make([]byte, 32)
	for i := range s {
		s[i] = hexString[rand.Intn(15)]
	}
	return string(s)
}

func getAddr(dataString string, base int) (string, string) {
	port, _ := strconv.ParseInt(dataString[base+12:base+16], 16, 64)
	data, _ := hex.DecodeString(dataString[base+16 : base+24])
	var ipAddr net.IP = data
	return ipAddr.String(), strconv.FormatInt(port, 10)
}

func stunTest(stunHost string, stunPort int, sendData string) map[string]string {
	val := map[string]string{
		"Resp": "0", "ExternalIP": "",
		"ExternalPort": "",
		"SourceIP":     "",
		"SourcePort":   "",
		"ChangedIP":    "",
		"ChangedPort":  ""}
	addr := new(net.UDPAddr)
	addr.Port = SourcePort
	conn, _ := net.ListenUDP("udp", addr)
	defer conn.Close()
	tranId := genTranId()
	strLen := fmt.Sprintf("%04v", len(sendData)/2)
	stunPortStr := strconv.Itoa(stunPort)
	stunAddr := stunHost + ":" + stunPortStr
	serverUDPAddr, _ := net.ResolveUDPAddr("udp", stunAddr)
	data, _ := hex.DecodeString(BindRequestMsg + strLen + tranId + sendData)
	buf := make([]byte, 2048)
	timeout := 100
	count := 3
	udpOffset := 0
	_, err := conn.WriteTo(data, serverUDPAddr)
	if err != nil {
		log.Println("WriteTo->")
		log.Println(err)
		val["Resp"] = "0"
		return val
	}
	//100ms for timeout
	conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Millisecond))
	for {
		offset, _, read_err := conn.ReadFromUDP(buf)
		udpOffset = offset
		if read_err == nil {
			break
		}
		if count > 0 {
			count = count - 1
		} else {
			log.Println("Due to count Resp->0")
			val["Resp"] = "0"
			return val
		}
	}
	recvData := hex.EncodeToString(buf[:udpOffset])
	tranidMatch := (recvData[8:40] == tranId) && (recvData[0:4] == BindResponseMsg)
	if tranidMatch {
		val["Resp"] = "1"
		messageLen := new(big.Int)
		messageLen.SetString(recvData[4:8], 16)
		base := 40
		attrType := recvData[base : base+4]
		//prot_fami := recv_data[base+8 : base+12]
		//attr_len default for 8byte and ip 4byte
		count := messageLen.Int64() / 12
		for count > 0 {
			attrType = recvData[base : base+4]
			attrLen, _ := strconv.ParseInt(recvData[base+4:base+8], 16, 64)
			if attrType == MappedAddress {
				ip, port := getAddr(recvData, base)
				val["ExternalIP"] = ip
				val["ExternalPort"] = port
				log.Println("MappedAddress", ip, ":", port)
			}
			if attrType == SourceAddress {
				ip, port := getAddr(recvData, base)
				val["SourceIP"] = ip
				val["SourcePort"] = port
				log.Println("SourceAddress", ip, ":", port)
			}
			if attrType == ChangedAddress {
				ip, port := getAddr(recvData, base)
				val["ChangedIP"] = ip
				val["ChangedPort"] = port
				log.Println("ChangedAddress", ip, ":", port)
			}
			//Two hex string for one byte
			base = base + (4+int(attrLen))*2
			count = count - 1
		}
	}
	log.Println("rawRecvData:", recvData)
	return val
}

func doNatTest(stunHost string, stunPort int, sourceIp string, sourcePort int) (string, map[string]string) {
	respVal := stunTest(stunHost, stunPort, "")
	//Do test1
	resp := respVal["Resp"]
	if resp != "1" {
		return "UDP Blocked", respVal
	}
	externalIp := respVal["ExternalIP"]
	externalPort := respVal["ExternalPort"]
	changedIp := respVal["ChangedIP"]
	changedPort := respVal["ChangedPort"]
	changedIntPort, _ := strconv.Atoi(changedPort)
	if externalIp == sourceIp {
		changeRequest := ChangeRequest + "0004" + "00000006"
		respVal = stunTest(stunHost, stunPort, changeRequest)
		resp := respVal["Resp"]
		if resp == "1" {
			return "Open Internet", respVal
		} else {
			return "Symmetric UDP Firewall", respVal
		}
	} else {
		changeRequest := ChangeRequest + "0004" + "00000006"
		respVal = stunTest(stunHost, stunPort, changeRequest)
		resp := respVal["Resp"]
		//Do test2
		if resp == "1" {
			return "Full Cone", respVal
		} else {
			respVal = stunTest(stunHost, stunPort, "")
			if respVal["Resp"] != "1" {
				return "Changed Address Error", respVal
			} else {
				if (externalIp == respVal["ExternalIP"]) && (externalPort == respVal["ExternalPort"]) {
					changePortRequest := ChangeRequest + "0004" + "00000002"
					//Do test3
					respVal = stunTest(changedIp, changedIntPort, changePortRequest)
					if respVal["Resp"] == "1" {
						return "Restric NAT", respVal
					} else {
						return "Restric Port NAT", respVal
					}
				} else {
					return "Symmetric NAT", respVal
				}
			}
		}

	}
}
func GetNatType(stunHost string, stunPort int) (string, map[string]string) {
	return doNatTest(stunHost, stunPort, SourceIP, SourcePort)
}
